"""Set of classes to implement a CTPRPC Server

You can quickly implement a server by simply instancing :class:`CTPRPCServer`
and then registering a class to be served

>>> class EchoServer:
...     def echo()
...         return 'HELO'
>>> srv = CTPRPCServer(('', 1110))
>>> srv.register_class(EchoServer)
>>> srv.serve()

"""

import socket, select, sys, traceback, threading, time, platform

if platform.system() != 'Windows':
    try:
        from multiprocessing import Process, Queue
    except:
        from threading import Thread as Process
        from Queue import Queue
else:
    from threading import Thread as Process
    from Queue import Queue

from CTPRPC3Common import *
from CTPRPCPickle import  CTPRPCSerializer

try:
    from tlslite.api import *
except:
    print 'System wide tlslite not found, falling back to internal one'
    from os3.net.externals.tlslite.api import *

try:
    test_sock_flags = socket.MSG_WAITALL
except:
    print 'System does not support MSG_WAITALL flag, trying to go without it'

class CTPRPCSingleKeepAlive(threading.Thread):
    """SingleKeepAlive objects take care of sending an HeartBeat packet to the
    client every 2 seconds. They will run inside their own separate thread
    and will run untile their CTPRPCSingleServer won't be stopped"""
    def __init__(self, parent):
        super(CTPRPCSingleKeepAlive, self).__init__()
        self.server = parent
        self.loop = True

    def dispatch_multiserver_request(self, message):
        if message:
            if message['type'] == 'fire_event':
                self.server.fire_event(message['event_name'], *message['args'], **message['kw'])
            elif message['type'] == 'quit':
                self.server.run = False
                self.server.retries = 0

    def run(self):
        while self.server.is_running() and self.loop:
            #this might fail if the main thread has already been closed,
            #it will simply raise an exception during interpreter shutdown
            try:
                self.server.send(HBeatPacket(self.server.status))
                up = True
            except:
                up = False

            if not up:
                continue

            try:
                #this can freeze for more than 2 seconds if the system is freezed
                #so it might happen that hbeats won't be sent for more than 2 seconds.
                message = self.server.communication_queue.get(True, 2)
            except:
                message = None
            self.dispatch_multiserver_request(message)

    def quit(self):
        self.loop = False

class CTPRPCSingleServer(Process):
    """:class:`CTPRPCSingleServer` takes care of handling every request of a specific client.
    
    Each client will be handled by their own :class:`CTPRPCSingleServer` instance.
    If multiprocessing module is available each SingleServer will run in
    a separate process, otherwise each :class:`CTPRPCSingleServer` will run in a separate thread"""
    def __init__(self, session, sock, cls=None, ** cls_args ):
        super(CTPRPCSingleServer, self).__init__()

        try:
            self.read_sock_flags = socket.MSG_WAITALL
        except:
            self.read_sock_flags = 0
        
        self.session = session
        self.sock = sock
        self.seq = 0
        self.recv_seq = 0
        self.obj = cls and cls( ** cls_args )
        if self.obj:
            self.obj.ctprpc = self
        self.events = {}
        self.dest = None
        self.send_lock = threading.Lock()
        self.events_lock = threading.Lock()
        self.serializer = CTPRPCSerializer()
        self.communication_queue = Queue()

        #: Status can be N for NOP, C for performing method CALL
        self.status = 'N'
        self.current_packet = None
        self.last_processed_call_request = None
        self.last_event_id = 0

        #Those are server timeouts, if nothing is received for more then self.timeout seconds
        #a soft timeout is caused, if more then self.retries soft timeout happen then a
        #real timeout will happen
        self.retries = 6
        self.timeout = 10

        #Those are to permit monkey patching
        self.select = select.select
        self.logger = CTPRPCServerEventsLogger()

    def send(self, pkt):
        """
        Sends a :class:`os3.net.CTPRPC3Common.CTPRPCPacket` to the client
        """
        self.send_lock.acquire()

        if not self.dest or not self.sock:
            self.send_lock.release()
            return

        pkt.seq = self.seq + 1
        pkt.session = self.session

        for pkt in pkt.generated_split():
            self.seq = pkt.seq
            if not reliable_send(self, pkt, self.dest):
                break

        self.send_lock.release()

    def fire_event(self, event_name, *args, **kw):
        """
        Fires the given event only to the client connected
        to this server
        """
        self.events_lock.acquire()
        if self.events.has_key(event_name):
            self.last_event_id += 1
            call_request = {'id':self.last_event_id,
                            'method':event_name,
                            'args':args,
                            'kw':kw}
            pkt = CallPacket(self.serializer, call_request)
            self.send(pkt)
        else:
            self.logger.warn_firing_unregistered_event(event_name)
        self.events_lock.release()

    def subscribe_event(self, event_name):
        """
        Registers the client as expecting to receive
        notifications for the given event
        """
        self.events_lock.acquire()
        self.events[event_name] = True
        self.events_lock.release()

    def read(self):
        """
        Will read a :class:`os3.net.CTPRPC3Common.CTPRPCPacket` from the network waiting
        for it at least 60 seconds. If no data is received in 60 seconds
        a :class:`CTPRPCError` exception is raised
        """
        timeouts = 0

        while timeouts < self.retries:
            try:
                rl, wl, xl = self.select ([self.sock], [], [self.sock], 10.0)
            except KeyboardInterrupt:
                break

            if xl or not rl:
                timeouts += 1
                continue

            if rl:
                self.send_lock.acquire()
                try:
                    data, addr = self.sock.recvfrom(1500, self.read_sock_flags|socket.MSG_DONTWAIT)
                except:
                    self.send_lock.release()
                    continue
                self.dest = addr

                timeouts = 0
                try:
                    pkt = CTPRPCPacket.load(data)
                except:
                    self.logger.err_malformed_packet(self.session, 'Corrupted Packet')
                    self.send_lock.release()
                    return None

                try:
                    self.sock.sendto('K-%s' % pkt.seq, self.dest)
                    if self.recv_seq < pkt.seq:
                        self.recv_seq = pkt.seq
                    else:
                        return None

                    if self.current_packet and not self.current_packet.marker:
                        self.current_packet.append(pkt)
                    else:
                        self.current_packet = pkt

                    if not self.current_packet.marker:
                        return None

                    if not self.current_packet.is_valid:
                        self.logger.err_malformed_packet(self.session, 'Invalid Packet')
                        self.current_packet = None
                        return None

                    self.current_packet.decompress()
                    return self.current_packet
                finally:
                    self.send_lock.release()
                    

        raise CTPRPCError('Hard timeout while waiting something from the client')

    def close(self):
        self.sock.close()
        self.sock = None

    def call_request_is_duplicate(self, p):
        return self.last_processed_call_request != None and p['id'] <= self.last_processed_call_request['id']

    def call_request_update(self, p, res):
        self. last_processed_call_request = {'id':p['id'], 'res':res}

    def call_request_answer(self):
        return AckPacket(self.serializer, self.last_processed_call_request['id'], self.last_processed_call_request['res'])

    def _recover_files(self, call):
        for file in self.serializer.files:
            filename, fileo = file
            fileo.write(call['files'][filename])
            fileo.seek(0)
        self.serializer.files = []

    def handle_packet(self, pkt):
        """
        Will parse and handle one incoming :class:`os3.net.CTPRPC3Common.CTPRPCPacket`
        """
        pkt = parse_packet(pkt)

        if isinstance(pkt, SubscribePacket):
            try:
                self.subscribe_event(pkt.event_name)
            except Exception, e:
                self.logger.err_malformed_packet(self.session, e)
                return
        elif isinstance(pkt, CallPacket):
            try:
                p = pkt.get_parsed_method(self.serializer)
            except Exception, e:
                self.logger.err_malformed_packet(self.session, e)
                return

            m = getattr(self.obj, p['method'], None)
            if not m:
                self.send(NotFoundPacket(self.serializer, p['id']))
            else:
                try:
                    if not self.call_request_is_duplicate(p):
                        self.status = 'C'
                        self._recover_files(p)
                        res = m(*p['args'], **p['kw'])
                        self.call_request_update(p, res)
                    self.status = 'N'
                    self.send(self.call_request_answer())
                except Exception, e:
                    tb = traceback.extract_tb(sys.exc_info()[2])
                    self.send(ExceptionPacket(self.serializer, p['id'], e.__class__, e, tb))

    def run(self):
        """
        Will wait and serve requests received from its :class:`os3.net.CTPRPC3Client.CTPRPCClient`
        
        Will also fork a :class:`CTPRPCSingleKeepAlive` instance to send
        :class:`os3.net.CTPRPC3Common.HBeatPacket` and handle requests from the master :class:`CTPRPCServer` instance
        """
        self.run = True
        keep_aliver = CTPRPCSingleKeepAlive(self)
        keep_aliver.start()

        while self.run:
            try:
                pkt = self.read()
                if pkt:
                    self.handle_packet(pkt)
            except CTPRPCError, e:
                self.run = self.logger.err_hard_timeout(self.session)
            except Exception, e:
                traceback.print_exception(*sys.exc_info())
                self.run = self.logger.err_malformed_packet(self.session, e)

        keep_aliver.quit()
        self.close()
        return False

    def is_running(self):
        try:
            return self.isAlive()
        except:
            return self.is_alive()

class CTPRPCServer(object):
    """
    Creates a CTPRPC server.
    if **multiprocessing** module is available the server will be created
    using multiple processes, one process for each client (Prefork model).
    If only the **threading** module is available the server will use threads
    instead. 
    
    The server will fork multiple :class:`CTPRPCSingleServer` instances to handle
    client requests, one instance for each client.
    
    The server requires an address and port to **listen** to in the form ('address', port).
    
    Also a **passwd_file** can be given, this will enable username and password authentication
    and will enable encryption. The file has to be in the form of a passwd file.
    """
    def __init__(self, listen, passwd_file=None):
        self.listen_address = listen[0]
        self.listen_port = listen[1]
        self.children = []

        self.server_sock = socket.socket()
        self.server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_sock.bind((self.listen_address, self.listen_port))
        self.server_sock.settimeout(2)

        self.registered_class = None
        self.registered_class_args = None
        self.session = 0

        #Those are for monkey patching
        self.logger = CTPRPCServerEventsLogger()
        self.select = select.select

        self._authdb = self._parse_auth_file(passwd_file)

    def _parse_auth_file(self, passwd_file):
        if not passwd_file:
            return None

        authdb = VerifierDB()
        try:
            pwd = file(passwd_file)
            for l in pwd:
                l = l.strip()
                if not l or l[0] == '#':
                    continue
                l = l.split (':')
                authdb[l[0].strip()] = VerifierDB.makeVerifier(l[0].strip(), l[1].strip(), 2048)
        except:
            self.logger.fatal_auth_file()
            authdb = None
        return authdb

    def register_class(self, cls, ** class_args):
        """
        Registers a class to be served

	You can specify all the parameters the class need when it will be instantiated.
	For example:

		register_class ( ClassName, param1 = 1, param2 = 2, ... )
        """
        self.registered_class = cls
	self.registered_class_args = class_args

    def register_event(self, event_name):
        """
        Registers the existance of an event.
        
        You actually don't need to call this method, it is here only for
        compatibility reasons, any event can be fired, if there is any client listening for it it will be dispatched.
        """
        return

    def fire_event(self, event_name, *args, **kw):
        """
        Fires an event to every children.
        
        Each client listening for the event will receive it and will handle it
        """
        for child in self.children:
            child.communication_queue.put({'type':'fire_event',
                                            'event_name':event_name,
                                            'args':args,
                                            'kw':kw})

    def serveOnce(self):
        """
        Servers one connection request, please use :meth:`serve` or :meth:`serve_forever`
        instead of calling this multiple times
        """
        try:
            rl, wl, xl = self.select ([self.server_sock], [], [], 60.0)
        except KeyboardInterrupt:
            print 'GOT Keyboard Interrupt, waiting for children to exit...'
            rl = []
            xl = None
            self.run = False

        if xl:
            self.logger.fatal_network_error()
            self.run = False

        for con in rl:
            cli, addr = con.accept()
            connection = None

            try:
                if self._authdb:
                    connection = TLSConnection(cli)
                    connection.closeSocket = True
                    connection.handshakeServer(verifierDB=self._authdb)
                else:
                    connection = cli
            except:
                connection.close()
                connection = None

            if connection:
                self.session += 1
                srv_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                srv_sock.bind((self.listen_address, 0))

                self.logger.accepted(cli.getpeername(), self.session)
                srv = CTPRPCSingleServer(self.session, srv_sock, self.registered_class, ** self.registered_class_args )
                srv.start()
                self.children.append(srv)

                cli_data = "%s:%s" % (self.session, srv_sock.getsockname()[1])
                try:
                    connection.send(cli_data)
                    connection.close()
                except:
                    pass

    def serve(self):
        """
        Will serve client connections requests, can also be called as :meth:`serve_forever` for
        compatibility reasons
        """
        self.server_sock.listen(10)
        self.run = True

        while self.run:
            self.serveOnce()
            self._clean_dead_children()

        self._quit_children()
    serve_forever = serve

    def _clean_dead_children(self):
        to_remove = []

        for child in self.children:
            if not child.is_running():
                to_remove.append(child)

        for child in to_remove:
            self.children.remove(child)

    def _quit_children(self):
        for child in self.children:
            child.communication_queue.put({'type':'quit'})

def main():
    import sys, time
    class Echo(object):
        def echo(self, *args):
            return args

        def ping(self):
            return 'pong'

    port = 8000
    srv = CTPRPCServer(('', port))
    srv.register_class(Echo)
    srv.serve()

if __name__ == '__main__':
    main()

